# -*- coding: utf-8 -*-
"""
下载器核心模块
提供视频下载、进度管理、多线程下载等功能
"""

import os
import time
import threading
from pathlib import Path
from typing import Dict, List, Optional, Callable, Tuple
from enum import Enum
from dataclasses import dataclass
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed

from .bilibili_api import BilibiliAPI
from .utils import (
    validate_bilibili_url,
    ensure_directory,
    get_safe_filename,
    format_file_size,
    format_speed,
    calculate_eta,
    setup_logger
)


class DownloadStatus(Enum):
    """下载状态枚举"""
    IDLE = "idle"                    # 空闲
    PREPARING = "preparing"          # 准备中
    DOWNLOADING = "downloading"      # 下载中
    PAUSED = "paused"               # 已暂停
    COMPLETED = "completed"         # 已完成
    FAILED = "failed"               # 失败
    CANCELLED = "cancelled"         # 已取消


@dataclass
class DownloadProgress:
    """下载进度信息"""
    total_size: int = 0              # 总大小（字节）
    downloaded_size: int = 0         # 已下载大小（字节）
    speed: float = 0.0              # 下载速度（字节/秒）
    eta: str = "--:--"              # 预计剩余时间
    percentage: float = 0.0         # 完成百分比
    status: DownloadStatus = DownloadStatus.IDLE
    current_file: str = ""          # 当前下载文件名
    error_message: str = ""         # 错误信息


class DownloadTask:
    """下载任务类"""
    
    def __init__(self, task_id: str, url: str, save_path: str, filename: str = ""):
        self.task_id = task_id
        self.url = url
        self.save_path = Path(save_path)
        self.filename = filename
        self.progress = DownloadProgress()
        self.video_info = None
        self.stream_info = None
        self.selected_quality = 80  # 默认1080P
        self.selected_page = 1      # 默认第一P
        
        # 下载控制
        self._stop_event = threading.Event()
        self._pause_event = threading.Event()
        self._download_thread = None
        
        # 回调函数
        self.progress_callback: Optional[Callable[[DownloadProgress], None]] = None
        self.completion_callback: Optional[Callable[[str, bool], None]] = None
        
        # 日志
        self.logger = setup_logger(f'DownloadTask-{task_id}')
    
    def set_progress_callback(self, callback: Callable[[DownloadProgress], None]):
        """设置进度回调函数"""
        self.progress_callback = callback
    
    def set_completion_callback(self, callback: Callable[[str, bool], None]):
        """设置完成回调函数"""
        self.completion_callback = callback
    
    def is_active(self) -> bool:
        """检查任务是否处于活动状态"""
        return self.progress.status in [DownloadStatus.PREPARING, DownloadStatus.DOWNLOADING]
    
    def is_paused(self) -> bool:
        """检查任务是否已暂停"""
        return self.progress.status == DownloadStatus.PAUSED
    
    def is_completed(self) -> bool:
        """检查任务是否已完成"""
        return self.progress.status == DownloadStatus.COMPLETED
    
    def stop(self):
        """停止下载任务"""
        self._stop_event.set()
        if self._download_thread and self._download_thread.is_alive():
            self._download_thread.join(timeout=5)
        self.progress.status = DownloadStatus.CANCELLED
        self.logger.info(f"任务已停止: {self.task_id}")
    
    def pause(self):
        """暂停下载任务"""
        if self.progress.status == DownloadStatus.DOWNLOADING:
            self._pause_event.set()
            self.progress.status = DownloadStatus.PAUSED
            self.logger.info(f"任务已暂停: {self.task_id}")
    
    def resume(self):
        """恢复下载任务"""
        if self.progress.status == DownloadStatus.PAUSED:
            self._pause_event.clear()
            self.progress.status = DownloadStatus.DOWNLOADING
            self.logger.info(f"任务已恢复: {self.task_id}")


class BilibiliDownloader:
    """B站视频下载器"""
    
    def __init__(self, max_concurrent_downloads: int = 3, chunk_size: int = 1024*1024):
        self.api = BilibiliAPI()
        self.max_concurrent_downloads = max_concurrent_downloads
        self.chunk_size = chunk_size  # 1MB chunks
        
        # 任务管理
        self.tasks: Dict[str, DownloadTask] = {}
        self.active_downloads = 0
        self._task_counter = 0
        self._lock = threading.Lock()
        
        # 线程池
        self.executor = ThreadPoolExecutor(max_workers=max_concurrent_downloads)
        
        # 日志
        self.logger = setup_logger('BilibiliDownloader')
        
        # 请求会话
        self.session = requests.Session()
        self.session.headers.update({
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36',
            'Referer': 'https://www.bilibili.com/'
        })
    
    def create_download_task(self, url: str, save_path: str, filename: str = "") -> Optional[str]:
        """
        创建下载任务
        
        Args:
            url: 视频URL
            save_path: 保存路径
            filename: 文件名（可选）
            
        Returns:
            str: 任务ID，失败返回None
        """
        if not validate_bilibili_url(url):
            self.logger.error(f"无效的B站URL: {url}")
            return None
        
        # 确保保存目录存在
        if not ensure_directory(save_path):
            self.logger.error(f"无法创建保存目录: {save_path}")
            return None
        
        # 生成任务ID
        with self._lock:
            self._task_counter += 1
            task_id = f"task_{self._task_counter}_{int(time.time())}"
        
        # 创建任务
        task = DownloadTask(task_id, url, save_path, filename)
        self.tasks[task_id] = task
        
        self.logger.info(f"创建下载任务: {task_id}")
        return task_id
    
    def start_download(self, task_id: str, quality: int = 80, page: int = 1) -> bool:
        """
        开始下载任务
        
        Args:
            task_id: 任务ID
            quality: 视频质量
            page: 视频分P
            
        Returns:
            bool: 是否成功开始
        """
        if task_id not in self.tasks:
            self.logger.error(f"任务不存在: {task_id}")
            return False
        
        task = self.tasks[task_id]
        if task.is_active():
            self.logger.warning(f"任务已在运行: {task_id}")
            return False
        
        # 检查并发下载限制
        with self._lock:
            if self.active_downloads >= self.max_concurrent_downloads:
                self.logger.warning(f"达到最大并发下载数限制: {self.max_concurrent_downloads}")
                return False
            self.active_downloads += 1
        
        # 设置任务参数
        task.selected_quality = quality
        task.selected_page = page
        task.progress.status = DownloadStatus.PREPARING
        
        # 启动下载线程
        task._download_thread = threading.Thread(
            target=self._download_worker,
            args=(task,),
            daemon=True
        )
        task._download_thread.start()
        
        self.logger.info(f"开始下载任务: {task_id}")
        return True
    
    def _download_worker(self, task: DownloadTask):
        """
        下载工作线程
        
        Args:
            task: 下载任务
        """
        try:
            # 解析视频信息
            if not self._prepare_download(task):
                return
            
            # 开始下载
            self._perform_download(task)
            
        except Exception as e:
            self.logger.error(f"下载任务异常: {task.task_id}, 错误: {e}")
            task.progress.status = DownloadStatus.FAILED
            task.progress.error_message = str(e)
        
        finally:
            # 减少活动下载计数
            with self._lock:
                self.active_downloads = max(0, self.active_downloads - 1)
            
            # 调用完成回调
            if task.completion_callback:
                success = task.progress.status == DownloadStatus.COMPLETED
                task.completion_callback(task.task_id, success)
    
    def _prepare_download(self, task: DownloadTask) -> bool:
        """
        准备下载：获取视频信息和流地址
        
        Args:
            task: 下载任务
            
        Returns:
            bool: 是否准备成功
        """
        # 解析URL
        url_info = self.api.parse_video_url(task.url)
        if not url_info:
            task.progress.status = DownloadStatus.FAILED
            task.progress.error_message = "URL解析失败"
            return False
        
        # 获取视频信息
        task.video_info = self.api.get_video_info(url_info['video_id'])
        if not task.video_info:
            task.progress.status = DownloadStatus.FAILED
            task.progress.error_message = "获取视频信息失败"
            return False
        
        # 检查分P
        if task.selected_page > len(task.video_info['pages']):
            task.selected_page = 1
        
        page_info = task.video_info['pages'][task.selected_page - 1]
        cid = page_info['cid']
        
        # 获取流信息
        task.stream_info = self.api.get_video_stream_info(
            url_info['video_id'], cid, task.selected_quality
        )
        if not task.stream_info:
            task.progress.status = DownloadStatus.FAILED
            task.progress.error_message = "获取视频流信息失败"
            return False
        
        # 生成文件名
        if not task.filename:
            title = get_safe_filename(task.video_info['title'])
            if len(task.video_info['pages']) > 1:
                part_name = get_safe_filename(page_info['part'])
                task.filename = f"{title}_P{task.selected_page}_{part_name}.mp4"
            else:
                task.filename = f"{title}.mp4"
        
        task.progress.current_file = task.filename
        self.logger.info(f"准备完成: {task.task_id}, 文件: {task.filename}")
        return True
    
    def _perform_download(self, task: DownloadTask):
        """
        执行下载
        
        Args:
            task: 下载任务
        """
        task.progress.status = DownloadStatus.DOWNLOADING
        
        # 获取下载URL
        video_streams = task.stream_info.get('video_streams', [])
        audio_streams = task.stream_info.get('audio_streams', [])
        
        if not video_streams:
            task.progress.status = DownloadStatus.FAILED
            task.progress.error_message = "没有可用的视频流"
            return
        
        # 选择最佳视频流
        video_stream = video_streams[0]  # 简化处理，选择第一个
        video_url = video_stream.get('base_url', '')
        
        if not video_url:
            task.progress.status = DownloadStatus.FAILED
            task.progress.error_message = "无法获取视频下载地址"
            return
        
        # 下载文件
        file_path = task.save_path / task.filename
        
        try:
            self._download_file(task, video_url, file_path)
            
            if not task._stop_event.is_set():
                task.progress.status = DownloadStatus.COMPLETED
                task.progress.percentage = 100.0
                self.logger.info(f"下载完成: {task.task_id}")
            
        except Exception as e:
            task.progress.status = DownloadStatus.FAILED
            task.progress.error_message = f"下载失败: {e}"
            self.logger.error(f"下载失败: {task.task_id}, 错误: {e}")
    
    def _download_file(self, task: DownloadTask, url: str, file_path: Path):
        """
        下载单个文件
        
        Args:
            task: 下载任务
            url: 下载URL
            file_path: 保存路径
        """
        # 获取文件大小
        try:
            head_response = self.session.head(url, timeout=10)
            total_size = int(head_response.headers.get('content-length', 0))
            task.progress.total_size = total_size
        except:
            total_size = 0
        
        # 检查是否支持断点续传
        resume_pos = 0
        if file_path.exists():
            resume_pos = file_path.stat().st_size
            if resume_pos >= total_size:
                # 文件已完整
                task.progress.downloaded_size = total_size
                return
        
        # 设置请求头
        headers = {}
        if resume_pos > 0:
            headers['Range'] = f'bytes={resume_pos}-'
            task.progress.downloaded_size = resume_pos
        
        # 开始下载
        start_time = time.time()
        last_update_time = start_time
        
        try:
            response = self.session.get(url, headers=headers, stream=True, timeout=30)
            response.raise_for_status()
            
            mode = 'ab' if resume_pos > 0 else 'wb'
            with open(file_path, mode) as f:
                for chunk in response.iter_content(chunk_size=self.chunk_size):
                    # 检查停止信号
                    if task._stop_event.is_set():
                        break
                    
                    # 检查暂停信号
                    while task._pause_event.is_set() and not task._stop_event.is_set():
                        time.sleep(0.1)
                    
                    if chunk:
                        f.write(chunk)
                        task.progress.downloaded_size += len(chunk)
                        
                        # 更新进度
                        current_time = time.time()
                        if current_time - last_update_time >= 0.5:  # 每0.5秒更新一次
                            self._update_progress(task, current_time - start_time)
                            last_update_time = current_time
            
            # 最终进度更新
            self._update_progress(task, time.time() - start_time)
            
        except requests.exceptions.RequestException as e:
            raise Exception(f"网络请求失败: {e}")
        except IOError as e:
            raise Exception(f"文件写入失败: {e}")
    
    def _update_progress(self, task: DownloadTask, elapsed_time: float):
        """
        更新下载进度
        
        Args:
            task: 下载任务
            elapsed_time: 已用时间（秒）
        """
        if task.progress.total_size > 0:
            task.progress.percentage = (task.progress.downloaded_size / task.progress.total_size) * 100
        
        if elapsed_time > 0:
            task.progress.speed = task.progress.downloaded_size / elapsed_time
            task.progress.eta = calculate_eta(
                task.progress.downloaded_size,
                task.progress.total_size,
                task.progress.speed
            )
        
        # 调用进度回调
        if task.progress_callback:
            task.progress_callback(task.progress)
    
    def pause_download(self, task_id: str) -> bool:
        """
        暂停下载
        
        Args:
            task_id: 任务ID
            
        Returns:
            bool: 是否成功暂停
        """
        if task_id not in self.tasks:
            return False
        
        task = self.tasks[task_id]
        task.pause()
        return True
    
    def resume_download(self, task_id: str) -> bool:
        """
        恢复下载
        
        Args:
            task_id: 任务ID
            
        Returns:
            bool: 是否成功恢复
        """
        if task_id not in self.tasks:
            return False
        
        task = self.tasks[task_id]
        task.resume()
        return True
    
    def stop_download(self, task_id: str) -> bool:
        """
        停止下载
        
        Args:
            task_id: 任务ID
            
        Returns:
            bool: 是否成功停止
        """
        if task_id not in self.tasks:
            return False
        
        task = self.tasks[task_id]
        task.stop()
        
        # 从任务列表中移除
        del self.tasks[task_id]
        return True
    
    def get_task_progress(self, task_id: str) -> Optional[DownloadProgress]:
        """
        获取任务进度
        
        Args:
            task_id: 任务ID
            
        Returns:
            DownloadProgress: 进度信息
        """
        if task_id not in self.tasks:
            return None
        
        return self.tasks[task_id].progress
    
    def get_all_tasks(self) -> Dict[str, DownloadProgress]:
        """
        获取所有任务的进度
        
        Returns:
            dict: 任务ID到进度信息的映射
        """
        return {task_id: task.progress for task_id, task in self.tasks.items()}
    
    def cleanup_completed_tasks(self):
        """
        清理已完成的任务
        """
        completed_tasks = [
            task_id for task_id, task in self.tasks.items()
            if task.progress.status in [DownloadStatus.COMPLETED, DownloadStatus.FAILED, DownloadStatus.CANCELLED]
        ]
        
        for task_id in completed_tasks:
            del self.tasks[task_id]
        
        self.logger.info(f"清理了 {len(completed_tasks)} 个已完成任务")
    
    def shutdown(self):
        """
        关闭下载器
        """
        # 停止所有任务
        for task in self.tasks.values():
            task.stop()
        
        # 关闭线程池
        self.executor.shutdown(wait=True)
        
        # 关闭API会话
        self.api.close()
        self.session.close()
        
        self.logger.info("下载器已关闭")